package com.smile.commons.tree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TreeUtil {

    /**
     * 将平铺的数据列表转换为树形结构
     * @param nodes 所有节点列表
     * @param rootParentId 根节点的父ID
     * @param <T> 节点ID类型
     * @return 树形结构列表
     */
    public static <T> List<TreeNode<T>> buildTree(List<TreeNode<T>> nodes, T rootParentId) {
        // 创建一个ID到节点的映射
        Map<T, TreeNode<T>> nodeMap = new HashMap<>();
        for (TreeNode<T> node : nodes) {
            nodeMap.put(node.getId(), node);
        }

        // 创建一个用于存储根节点的列表
        List<TreeNode<T>> rootNodes = new ArrayList<>();

        // 遍历所有节点，将节点挂载到其父节点上
        for (TreeNode<T> node : nodes) {
            T parentId = node.getParentId();
            if (rootParentId.equals(parentId)) {
                rootNodes.add(node);
            } else {
                TreeNode<T> parentNode = nodeMap.get(parentId);
                if (parentNode != null) {
                    List<TreeNode<T>> children = parentNode.getChildren();
                    if (children == null) {
                        children = new ArrayList<>();
                        parentNode.setChildren(children);
                    }
                    children.add(node);
                }
            }
        }
        return rootNodes;
    }

    /**
     * 从平铺的数据列表中找到所有根节点
     * @param nodes 所有节点列表
     * @param rootParentId 根节点的父ID
     * @param <T> 节点ID类型
     * @return 根节点列表
     */
    public static <T> List<TreeNode<T>> findRootNodes(List<TreeNode<T>> nodes, T rootParentId) {
        List<TreeNode<T>> rootNodes = new ArrayList<>();
        for (TreeNode<T> node : nodes) {
            if (rootParentId.equals(node.getParentId())) {
                rootNodes.add(node);
            }
        }
        return rootNodes;
    }
}

